package net.jcip.examples.chapter7;

import java.util.concurrent.*;

public class TimedRunFix {

    private final ExecutorService taskExec;


    public TimedRunFix(ExecutorService taskExec) {
        this.taskExec = taskExec;
    }

    public void timedRun(Runnable r, long timeout, TimeUnit unit) throws InterruptedException {
            Future<?> task = taskExec.submit(r);
            try {
                task.get(timeout, unit);
            } catch (TimeoutException e) {
                // task will be cancelled below
            } catch (ExecutionException e) {
                // exception thrown in task; rethrow
                throw launderThrowable(e.getCause());
            } finally {
                // Harmless if task already completed
                task.cancel(true); // interrupt if running
            }
        }

    private static RuntimeException launderThrowable(Throwable t) {
        if (t instanceof RuntimeException) return (RuntimeException) t;
        else if (t instanceof Error) throw (Error) t;
        else throw new IllegalStateException("Not unchecked", t);
    }

    public void shutdown() {
        taskExec.shutdown();
    }

}
